/*  Copyright (c) 2011, Philip Busch <vzxwco@gmail.com>
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *
 *   - Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *   - Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 *  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 */


#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <unistd.h>
#include <fcntl.h>
#include <assert.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>
#include "list.h"
#include "stream.h"
#include "nc.h"


/* the main data structure that all these functions operate on */
#define NC_IDSIZE 256
#define NC_ERRSIZE 1024
#define NC_HOSTSIZE 128
#define NC_PORTSIZE 16
#define NC_ADDRSIZE INET6_ADDRSTRLEN
#define NC_CHUNKSIZE 8192

struct nc {
	char	 id[NC_IDSIZE];        /* connection ID */
	char	 host[NC_HOSTSIZE];    /* hostname, as supplied by the user */
	char     addr[NC_ADDRSIZE];    /* target IP address */
	char	 port[NC_PORTSIZE];    /* target port */
	char     errmsg[NC_ERRSIZE];   /* last error in human readable form */
	int	 sockfd;               /* our socket descriptor */
	int      type;                 /* client or server connection? */
	stream	*in;                   /* buffer for incoming data */
	stream	*out;                  /* buffer for outgoing data */
	size_t	 nsent;                /* # of bytes sent */
	size_t	 nrcvd;                /* # of bytes received */
	size_t	 ncallbacks;           /* # of callbacks */
	nc_recognizer	*recognizers;  /* array of pointers to recognizer functions */
	nc_handler	*handlers;     /* array of pointers to handler functions */
};



static node_l *connections;

nc *nc_create(int type)
{
	nc *n;

	n = calloc(1, sizeof(nc));
	if (n == NULL)
		return NULL;

	n->in = stream_create();
	if (n->in == NULL) {
		free (n);
		return NULL;
	}

	n->out = stream_create();
	if (n->out == NULL) {
		stream_destroy(n->in);
		free (n);
		return NULL;
	}

	n->sockfd = -1;
	n->type = type;

	// my aunt has a toaster where NULL-pointers aren't all zeros
	n->recognizers = NULL;
	n->handlers = NULL;

	return n;	
}

int nc_register_callback(nc *n, nc_recognizer r, nc_handler h)
{
	nc_recognizer *tmp_r;
	nc_handler *tmp_h;

	tmp_r = realloc(n->recognizers, (n->ncallbacks + 1) * sizeof(nc_recognizer));
	if (tmp_r == NULL)
		return -1;

	tmp_h = realloc(n->handlers, (n->ncallbacks + 1) * sizeof(nc_handler));
	if (tmp_h == NULL)
		return -1;

	tmp_r[n->ncallbacks] = r;
	tmp_h[n->ncallbacks] = h;

	n->recognizers = tmp_r;
	n->handlers = tmp_h;
	n->ncallbacks += 1;

	return 0;
}

static int nc_call(nc *n, nc_recognizer r)
{
	size_t i;

	for (i = 0; i < n->ncallbacks; i++)
		if (n->recognizers[i] == r)
			return (n->handlers[i](n, 0));

	return -1;
}

int nc_is_connect(nc *n)
{
	if (n) {}
	return 0;
}

int nc_is_disconnect(nc *n)
{
	if (n) {}
	return 0;
}

int nc_is_error(nc *n)
{
	if (n) {}
	return 0;
}

static void nc_set_strerror(nc *n, char *fmt, ...)
{
	va_list argp;

	va_start(argp, fmt);
	vsnprintf(n->errmsg, NC_ERRSIZE, fmt, argp);
	va_end(argp);

	n->errmsg[NC_ERRSIZE - 1] = '\0';
}

char *nc_strerror(nc *n)
{
	return n->errmsg;
}

static int nc_connect(nc *n, int sockfd, const struct sockaddr *addr, socklen_t addr_len, int nsec)
{
	int status, flags, error;
	socklen_t len;
	fd_set rset, wset;
	struct timeval tval;

	flags = fcntl(sockfd, F_GETFL, 0);
	fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);

	error = 0;
	status = connect(sockfd, addr, addr_len);
	if (status < 0)
		if (errno != EINPROGRESS) {
			nc_set_strerror(n, "connect(): %s", strerror(errno));
			return -1;
		}

	if (status == 0)
		goto done;
	
	FD_ZERO(&rset);
	FD_SET(sockfd, &rset);
	wset = rset;

	tval.tv_sec = nsec;
	tval.tv_usec = 0;

	status = select(sockfd+1, &rset, &wset, NULL, nsec ? &tval : NULL);
	if (status == 0) {
		nc_set_strerror(n, "connect(): %s", strerror(ETIMEDOUT));
		return -1;
	}

	if (FD_ISSET(sockfd, &rset) || FD_ISSET(sockfd, &wset)) {
		len = sizeof(error);
		if (getsockopt(sockfd, SOL_SOCKET, SO_ERROR, &error, &len) < 0) {
			nc_set_strerror(n, "select(): Solaris pending error");
			return -1;
		}
	} else {
		nc_set_strerror(n, "select(): sockfd not set");
		return -1;
	}

done:
	fcntl(sockfd, F_SETFL, flags);

	if (error) {
		nc_set_strerror(n, "nc_connect(): %s", strerror(error));
		return -1;
	}

	return 0;
}

static void nc_rawip2str(nc *n, struct addrinfo *res)
{
	void *addr;

	if (res->ai_family == AF_INET) {
		struct sockaddr_in *ipv4 = (struct sockaddr_in *)res->ai_addr;
		addr = &(ipv4->sin_addr);
	} else { // IPv6
		struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)res->ai_addr;
		addr = &(ipv6->sin6_addr);
	}

	inet_ntop(res->ai_family, addr, n->addr, NC_ADDRSIZE);
}

// get sockaddr, IPv4 or IPv6:
static void *get_in_addr(struct sockaddr *sa)
{
	if (sa->sa_family == AF_INET)
		return &(((struct sockaddr_in*)sa)->sin_addr);

	return &(((struct sockaddr_in6*)sa)->sin6_addr);
}

static nc *nc_accept(nc *n)
{
	struct sockaddr_storage their_addr;
	socklen_t sin_size;
	nc *client;
	size_t i;

	client = nc_create(NC_CLIENT);
	if (client == NULL)
		return NULL;

	client->sockfd = accept(n->sockfd, (struct sockaddr *)&their_addr, &sin_size);
	if (client->sockfd < 0) {
		// FIXME: handle EAGAIN
		free (client);
		nc_set_strerror(n, "accept(): %s", strerror(errno));
		return NULL;
	}

	inet_ntop(their_addr.ss_family, get_in_addr((struct sockaddr *)&their_addr), client->host, sizeof (client->host));
	printf("got connection from %s\n", client->host);

/*
	client->ncallbacks = n->ncallbacks;
	client->recognizers = n->recognizers;
	client->handlers = n->handlers;
*/

	// FIXME: wtf??

	for (i = 0; i < n->ncallbacks; i++) {
		nc_register_callback(client, n->recognizers[i], n->handlers[i]);
	}

	client->ncallbacks = n->ncallbacks;

	nc_call(client, &nc_is_connect);

	return client;
}


int nc_dial(nc *n, char *addr, char *port, short family, int nsec)
{
	int status;
	int yes = 1;
	struct addrinfo hints, *res, *p;

	if (addr == NULL) {
		if (n->type != NC_SERVER) {
			nc_set_strerror(n, "nc_dial(): no host address specified");
			return -1;
		} else {
			strncpy(n->host, "localhost", NC_HOSTSIZE);
		}
	} else strncpy(n->host, addr, NC_HOSTSIZE);

	if (port == NULL) {
		nc_set_strerror(n, "nc_dial(): no port specified");
		return -1;
	} else strncpy(n->port, port, NC_PORTSIZE);

	memset(&hints, 0, sizeof hints);
	hints.ai_family = family;
	hints.ai_socktype = SOCK_STREAM;
	if (addr == NULL)
		hints.ai_flags = AI_PASSIVE;

	status = getaddrinfo(addr, port, &hints, &res);
	if (status < 0) {
		nc_set_strerror(n, "getaddrinfo(): %s (bad host/port?)", gai_strerror(status));
		return -1;
	}

	// loop through the results and connect to the first we can
	for (p = res; p != NULL; p = p->ai_next) {
		n->sockfd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
		if (n->sockfd < 0) {
			nc_set_strerror(n, "socket(): %s", strerror(errno));
			return -1;
		}

		if (n->type == NC_SERVER) {
			if (setsockopt(n->sockfd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int)) < 0) {
				while(close(n->sockfd));
				n->sockfd = -1;
				nc_set_strerror(n, "setsockopt(): %s", strerror(errno));
				return -1;
			}

			if (bind(n->sockfd, p->ai_addr, p->ai_addrlen) < 0) {
				while(close(n->sockfd));
				continue;
			}

			if (listen(n->sockfd, NC_BACKLOG) < 0) {
				while(close(n->sockfd));
				n->sockfd = -1;
				nc_set_strerror(n, "listen(): %s", strerror(errno));
				return -1;
			}

			break;
		} else {  // n->type == NC_CLIENT
			status = nc_connect(n, n->sockfd, p->ai_addr, p->ai_addrlen, nsec);
			if (status < 0) {
				while(close(n->sockfd));
				n->sockfd = -1;
				continue;
			}

			break;
		}
	}

	if (p == NULL) {
		nc_set_strerror(n, "nc_dial(): failed to connect");
		return -1;
	}

	nc_rawip2str(n, p);
	freeaddrinfo(res);

	status = list_append(&connections, n);
	if (status < 0) {  // out of memory :(
		nc_set_strerror(n, "nc_dial(): %s", strerror(errno));
		while(close(n->sockfd));
		n->sockfd = -1;
		return -1;
	}

	// finally, call connect handler
	if (n->type == NC_CLIENT)
		nc_call(n, &nc_is_connect);

	return 0;
}

int nc_send(nc *n, char *buf, size_t len)
{
	return stream_write(n->out, buf, len);
}

int nc_recv(nc *n, char *buf, size_t len)
{
	return stream_read(n->in, buf, len);
}

int nc_peak(nc *n, char *buf, size_t len)
{
	return stream_peak(n->in, buf, len);
}

size_t nc_size(nc *n)
{
	return stream_size(n->in);
}

int nc_do_send(nc *n)
{
	char buf[NC_CHUNKSIZE];
	int nb_read, nb_sent;

	nb_read = stream_peak(n->out, buf, NC_CHUNKSIZE);
	if (nb_read == 0)
		return 0;

	nb_sent = send(n->sockfd, buf, nb_read, 0);
	if (nb_sent < 0) {
		if (errno == EINTR) {
			return 0;
		} else {
			nc_set_strerror(n, "nc_do_send(): %s", strerror(errno));
			return -1;
		}
	}

	// remove from outbuffer what we actually sent
	(void) stream_read(n->out, buf, nb_sent);

	n->nsent += nb_sent;

	return nb_sent;
}

int nc_do_recv(nc *n)
{
	char buf[NC_CHUNKSIZE];
	int nb_rcvd;

	nb_rcvd = recv(n->sockfd, buf, NC_CHUNKSIZE, 0);
	if (nb_rcvd < 0) {
		if (errno == EINTR) {
			return 0;
		} else {
			nc_set_strerror(n, "nc_do_recv(): %s", strerror(errno));
			return -1;
		}
	} else if (nb_rcvd > 0) {
		stream_write(n->in, buf, nb_rcvd);
		n->nrcvd += nb_rcvd;
	} else {  // disconnected, mark for removal
		n->sockfd = -1;
	}

	return nb_rcvd;
}

static int nc_remove_connection(nc *n)
{
	node_l *np;

	np = list_search(&connections, n);

	// np can never be NULL here, but let's be cautious
	assert(np != NULL);

	list_unlink(&connections, np);

	return 0;
}

static int nc_build_fdsets(fd_set *readfds, fd_set *writefds, fd_set *errfds)
{
	node_l *np;
	nc *conn;

	int maxfd = -1;

	np = list_get_first_node(&connections);

	while (np != NULL) {
		conn = list_get_data(&np);

		if (stream_size(conn->out)) {
			FD_SET(conn->sockfd, writefds);
		}

		FD_SET(conn->sockfd, readfds);
		FD_SET(conn->sockfd, errfds);

		maxfd = conn->sockfd > maxfd ? conn->sockfd : maxfd;

		np = list_get_next_node(&connections, np);
	}

	return maxfd;
}

static nc *nc_handle_fdsets(fd_set *readfds, fd_set *writefds, fd_set *errfds)
{
	node_l *np;
	nc *conn, *client;
	int nbytes;

	np = list_get_first_node(&connections);

	for (; np != NULL; np = list_get_next_node(&connections, np)) {
		conn = list_get_data(&np);

		if (FD_ISSET(conn->sockfd, errfds)) {
			return conn;
		}

		if (FD_ISSET(conn->sockfd, readfds)) {
			if (conn->type == NC_SERVER) {
				if ((client = nc_accept(conn)) != NULL)
					if (list_append(&connections, client) < 0)
						nc_hangup(client);

			} else {  // conn->type == NC_CLIENT
				nbytes = nc_do_recv(conn);
				switch (nbytes) {
					case -1:
						// call error handler
						if (nc_call(conn, &nc_is_error) < 0)
							return conn;
						break;
					case 0:
						// call disconnect handler
						if (nc_call(conn, &nc_is_disconnect) < 0)
							return conn;
						break;
					default:
						break;
				}
			}
		}

		if (FD_ISSET(conn->sockfd, writefds)) {
			nbytes = nc_do_send(conn);
			if (nbytes < 0) {
				// call error handler
				if (nc_call(conn, &nc_is_error) < 0)
					return conn;
			}
		}
	}

	return NULL;
}

static nc *nc_handle_connections()
{
	int ntoksz, status;
	size_t i;
	node_l *np;
	nc *conn;

	np = list_get_first_node(&connections);

	while (np != NULL) {
		conn = list_get_data(&np);

		// skip if instream is empty
		if (stream_size(conn->in) == 0) {
			np = list_get_next_node(&connections, np);
			continue;
		}

		for (i = 0; i < conn->ncallbacks; i++) {
			// call recognizer
			ntoksz = conn->recognizers[i](conn);
			if (ntoksz < 0) {
				return conn;
			}

			if (ntoksz > 0) {
				// call corresponding handler
				status = conn->handlers[i](conn, ntoksz);
				if (status < 0) {
					return conn;
				}

				// there may be more data available, start again
				if (stream_size(conn->in))
					i = -1;
			}

		}

		np = list_get_next_node(&connections, np);
	}

	return NULL;
}

nc *nc_talk()
{
	fd_set readfds, writefds, errfds;
	int maxfd, status;
	nc *conn;
	struct timeval tv;

	while (1) {
		do {
			FD_ZERO(&readfds);
			FD_ZERO(&writefds);
			FD_ZERO(&errfds);

			maxfd = nc_build_fdsets(&readfds, &writefds, &errfds);
			if (maxfd < 0)
				return NULL;

			tv.tv_sec = 1;
			tv.tv_usec = 0;

			// FIXME: delete me
			printf("\n");
			//nc_print(connections[0].data);
			nc_print_all();
		} while((status = select(maxfd+1, &readfds, &writefds, &errfds, &tv)) == 0);
		

		if (status < 0) {
			fprintf(stderr, "libnc: Philip said this cannot possibly happen, but it did...\n");
			exit(EXIT_FAILURE);
		}

		conn = nc_handle_fdsets(&readfds, &writefds, &errfds);
		if (conn != NULL)
			return conn;

		conn = nc_handle_connections();
		if (conn != NULL)
			return conn;
	}

	// we never reach this line
	return NULL;
}

void nc_hangup(nc *n)
{
	if (n == NULL)
		return;

	nc_remove_connection(n);

	stream_destroy(n->in);
	stream_destroy(n->out);
	free(n->recognizers);
	free(n->handlers);
	free(n);

	n = NULL;
}

void nc_print(nc *n)
{
	printf("ID:                %s\n", n->id);
	printf("host:              %s\n", n->host);
	printf("addr:              %s\n", n->addr);
	printf("port:              %s\n", n->port);
	printf("socket:            %d\n", n->sockfd);
	printf("# bytes sent:      %d\n", n->nsent);
	printf("# bytes rcvd:      %d\n", n->nrcvd);
	printf("# of callacks:     %d\n", n->ncallbacks);
	printf("size of inbuffer:  %d\n", stream_size(n->in));
	printf("size of outbuffer: %d\n", stream_size(n->out));
	printf("last error:        %s\n", n->errmsg[0]=='\0'?"none":n->errmsg);
}

void nc_print_all(void)
{
	node_l *conn;
	nc *n;

	conn = list_get_first_node(&connections);

	printf("=========================================================\n");

	while (conn != NULL) {
		n = conn->data;
		nc_print(n);

		printf("------------------------------------------------------\n");

		conn = list_get_next_node(&connections, conn);
	}
}

